package com.yanggu.flink.datastream_api.state.keyedstate

import cn.hutool.core.date.DateUtil
import com.yanggu.flink.datastream_api.pojo.Event
import com.yanggu.flink.datastream_api.source.ClickSource
import org.apache.flink.api.common.eventtime.{SerializableTimestampAssigner, WatermarkStrategy}
import org.apache.flink.api.common.state.{MapState, MapStateDescriptor}
import org.apache.flink.configuration.Configuration
import org.apache.flink.streaming.api.functions.KeyedProcessFunction
import org.apache.flink.streaming.api.scala._
import org.apache.flink.util.Collector

import java.util.Date

/**
 * 使用MapState来模拟一下事件时间滚动窗口
 */
object MapStateDemo {

  def main(args: Array[String]): Unit = {
    val env = StreamExecutionEnvironment.getExecutionEnvironment
    env.setParallelism(1)

    env
      .addSource(new ClickSource)
      .assignTimestampsAndWatermarks(WatermarkStrategy
        .forMonotonousTimestamps()
        .withTimestampAssigner(new SerializableTimestampAssigner[Event] {
          override def extractTimestamp(element: Event, recordTimestamp: Long) = element.getTimestamp
        })
      )
      .keyBy(data => data.getUrl)
      //10s的滚动窗口
      .process(new FakeWindowResult(10000L))
      .print()

    env.execute()

  }

}

class FakeWindowResult(windowSize: Long) extends KeyedProcessFunction[String, Event, String] {

  var windowMapState: MapState[Long, Long] = _

  override def open(parameters: Configuration): Unit = {
    val windowMapStateDescriptor = new MapStateDescriptor[Long, Long]("windowMapState", classOf[Long], classOf[Long])
    windowMapState = getRuntimeContext.getMapState(windowMapStateDescriptor)
  }

  override def processElement(value: Event, ctx: KeyedProcessFunction[String, Event, String]#Context, out: Collector[String]): Unit = {
    val windowStart = value.getTimestamp / windowSize * windowSize
    val windowEnd = windowStart + windowSize

    //更新窗口中的pv数据
    windowMapState.put(windowStart, windowMapState.get(windowStart) + 1)

    //注册窗口结束时间定时器, 触发窗口的输出
    //因为Watermark默认 maxTimestamp - 1
    ctx.timerService().registerEventTimeTimer(windowEnd - 1)

  }

  override def onTimer(timestamp: Long, ctx: KeyedProcessFunction[String, Event, String]#OnTimerContext, out: Collector[String]): Unit = {
    val windowEnd = timestamp + 1
    val windowStart = windowEnd - windowSize

    val pv = windowMapState.get(windowStart)
    out.collect(s"url: ${ctx.getCurrentKey}, 窗口开始时间: ${DateUtil.formatDateTime(new Date(windowStart))}, " +
      s"窗口结束时间: ${DateUtil.formatDateTime(new Date(windowEnd))}, pv: $pv")

    //清除状态中, 指定窗口的数据
    windowMapState.remove(windowStart)

  }

}
